################################################################################
#       Copyright (C) 2010 Michael Yurko <myurko@gmail.com>
#
#This file is part of pyQAP.
#
#pyQAP is free software: you can redistribute it and/or modify
#it under the terms of the GNU General Public License as published by
#the Free Software Foundation, either version 2 of the License, or
#(at your option) any later version.
#
#pyQAP is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#GNU General Public License for more details.
#
#You should have received a copy of the GNU General Public License
#along with pyQAP.  If not, see <http://www.gnu.org/licenses/>.
################################################################################

import numpy as np
import classes
import heuristic_helpers as hh
from math import log, exp
import cost
import vars
'''
Contains heuristics for the QAP.
'''

def first(problem):
    '''
    A test heuristic which returns the first solution found.

    Returns the partial solution that refers to the solution array
    [0, 1, 2, ... , n-1].

    :param problem: the problem from which to generate the partial solution
    :type problem: Problem

    Examples:

    >>> from pyQAP import *
    >>> problem = problems.QAPLIB('nug30')
    >>> psol = heuristics.first(problem)
    >>> psol.cost(problem)
    8060.0
    >>> psol
    Partial QAP solution: [ 0  1  2  3  4  5  ... 26 27 28 29] with 30 placed.
     Placed:30
     Unplaced:[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
    >>> problem = problems.QAPLIB('had16')
    >>> psol = heuristics.first(problem)
    >>> psol.cost(problem)
    4144.0
    >>> psol
    Partial QAP solution: [ 0  1  2  3 ... 13 14 15] with 16 placed.
     Placed:16
     Unplaced:[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
    '''

    psol = classes.PSolution()
    psol.sol = np.arange(problem.n)
    psol.unplaced = np.zeros((problem.n), dtype=np.int)
    psol.placed = problem.n
    psol.n = problem.n
    return psol

def naive_greedy(problem):
    '''
    A trivial heuristic that returns the solution using a simple greedy paring
    of costs and distances.

    Return a heuristic solution to problem by greedily pairing the highest flow
    with the lowest cost left. If n is odd, then the final facility is placed
    in the last remaining location.

    :param problem: the problem from which to generate the partial solution
    :type problem: Problem
    
    Examples:
    
    >>> from pyQAP import *
    >>> prob = problems.QAPLIB('nug12')
    >>> naive_greedy(prob)
    Partial QAP solution: [ 0  3  1  8  7 11  6  2  4  5  9 10] with 12 placed.
     Placed:12
     Unplaced:[0 0 0 0 0 0 0 0 0 0 0 0]
    >>> prob = problems.QAPLIB('had16')
    >>> naive_greedy(prob)
    Partial QAP solution: [ 0  3  4 10  5  1  7  8 15 13  2  6 12  9 11 14] with 16 placed.
     Placed:16
     Unplaced:[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
    >>> prob = problems.QAPLIB('esc32a')
    >>> naive_greedy(prob)
    Partial QAP solution: [ 0 12  1 13  2  3 14 17 20 22  7 15 10 18 11 25 23 26  4  5 16 24  6  8  9
     19 21 27 28 29 30 31] with 32 placed.
     Placed:32
     Unplaced:[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
    >>> prob = problems.QAPLIB('sko100f')
    >>> naive_greedy(prob)
    Partial QAP solution: [ 0  9  1 99  8 89 90 10 19 80  2 91  7 11 18 70 20 81 29 98 79  3 88 97 69
      6 12 92 17 21 71 78 28 87 82 30 60 39  4  5 93 13 96 68 16 61 59 50 22 27
     31 38 77 40 72 49 83 94 95 86 14 58 15 51 23 26 62 67 32 76 37 41 48 24 25
     52 85 84 33 57 66 36 42 63 75 47 34 35 74 56 44 53 45 73 43 65 46 64 54 55] with 100 placed.
     Placed:100
     Unplaced:[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
     0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
     0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
    '''

    #read the flow matrix into a list to sort by highest
    flow_list = []
    for i in xrange(problem.n):
        for j in xrange(i+1,problem.n):
            flow_list.append((problem.f[i,j],i,j))

    #sort the flows with the highest first
    get_key = lambda t: t[0]
    flow_list.sort(key=get_key,reverse=True)

    #read the distance matrix into a list to sort by highest
    distance_list = []
    for i in xrange(problem.n):
        for j in xrange(i+1,problem.n):
            distance_list.append((problem.d[i,j],i,j))

    #sort the distances with the lowest first
    distance_list.sort(key=get_key)

    #initiate sets of unused locations and facilities
    unused_locations = set(range(problem.n))
    unused_facilities = set(range(problem.n))

    #initate solution array to return
    sol_array = np.empty(problem.n, dtype=np.int)

    #get list iterators
    flow_iter = flow_list.__iter__()
    distance_iter = distance_list.__iter__()

    #loop until no more than one facility is unassigned
    while len(unused_facilities) > 1:

        #get highest remaining flow
        hflow = flow_iter.next()
        while (hflow[1] not in unused_facilities) or (hflow[2] not in unused_facilities):
            hflow = flow_iter.next()

        #get lowest remaining distance
        ldist = distance_iter.next()
        while (ldist[1] not in unused_locations) or (ldist[2] not in unused_locations):
            ldist = distance_iter.next()

        #assign highest remaining flow to lowest remaining distance
        sol_array[ldist[1]] = hflow[1]
        sol_array[ldist[2]] = hflow[2]

        #remove used facilities and locations from set of unused
        unused_facilities.remove(hflow[1])
        unused_facilities.remove(hflow[2])
        unused_locations.remove(ldist[1])
        unused_locations.remove(ldist[2])

    #if n is odd assign the remaining location and facility
    if problem.n%2==1:
        last_loc = unused_locations.pop()
        last_fac = unused_facilities.pop()
        sol_array[last_loc] = last_fac

    #return PSolution of sol_array
    psol = classes.PSolution()
    psol.placed = problem.n
    psol.unplaced = np.zeros((problem.n), dtype=np.int)
    psol.sol = sol_array
    psol.n = problem.n
    return psol

def gradient_projection(problem, initial = None, max_iter = 100):
    '''
    Uses a gradient projection based heuristic to return a heuristic solution to
    to the problem.

    Return a heuristic solution to the given problem instance
    using a gradient projection algorithm starting with the given
    initial solution.

    :param problem: the problem from which to generate the partial solution
    :type problem: Problem
    
    Examples:
    
    >>> from pyQAP import *
    >>> prob = problems.QAPLIB('nug12')
    >>> vars.verbose = False
    >>> gradient_projection(prob)
    11 Improvements. Overall 4.882435
    Partial QAP solution: [ 0  1  2  3  4  5  6  7  8  9 10 11] with 12 placed.
     Placed:12
     Unplaced:[0 0 0 0 0 0 0 0 0 0 0 0]
    >>> prob = problems.QAPLIB('had16')
    >>> gradient_projection(prob)
    11 Improvements. Overall 13.417619
    Partial QAP solution: [ 0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15] with 16 placed.
     Placed:16
     Unplaced:[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
    >>> prob = problems.QAPLIB('tai12a')
    >>> gradient_projection(prob)
    11 Improvements. Overall 7185.001692
    Partial QAP solution: [ 0  1  2  3  4  5  6  7  8  9 10 11] with 12 placed.
     Placed:12
     Unplaced:[0 0 0 0 0 0 0 0 0 0 0 0]
    '''

    #unpack vars
    n = problem.n
    f = problem.f
    d = problem.d

    #main body of BestP
    if initial == None:
        x = np.asmatrix(np.identity(n))      #start with identity matrix
    else:
        x = np.matrix(initial)
    best_p = x.copy()             #keep track of best x so far
    itr = 0                 #keep track of iteration
    v = hh.qapobj(x,problem)   #get first cost
    best_v = v              #keep track of best v so far
    v0 = v                  #store original cost
    nb=0                    #store number of improvements
    olcv = True             #outer loop control variable

    while olcv:
        p_old = hh.round_to_perm(x)
        direction = -(f.T*x*d+f*x*d.T) #neg gradient = search dir
        #do a line search
        beta = .9
        alpha = 1/beta
        ilcv = True                     #inner loop control variable

        while ilcv:
            alpha = alpha * beta
            vplus = hh.qapobj(x+alpha*direction,problem)
            ilcv = vplus>v

        #new point
        x = x+alpha*direction

        #sanity check
        if vplus>v:
            print "No descent"

        #project back onto feasible space
        x = hh.proj_o(x,itr%2)             #non-convex set
        ilcv = True
        while ilcv:
            y = x.copy()
            if itr == 0:
                x, a, q, r, b = hh.proj_e(x)
            else:
                x = hh.proj_e(x, a, q, r, b)[0]
            x = hh.proj_n(x)
            if np.linalg.norm(y-x) < 1e-6:
                ilcv = False

        #look at potential improvements
        itr += 1
        v = hh.qapobj(x, problem)
        p = hh.round_to_perm(x)
        ve = hh.qapobj(p, problem)
        if ve < best_v:
            best_v = v
            best_p = p
            nb += 1
        if ((itr>max_iter)or (itr>10 and np.linalg.norm(p-p_old) < 1e-5)):
            olcv = False
    if vars.verbose:
        print "%d Improvements. Overall %f" % (nb, v0-best_v)
    #return a partial solution instance
    sol = hh.convert_x_to_array(best_p)
    unplaced = np.zeros((problem.n), dtype=np.int)
    psol = classes.PSolution(problem.n, None, sol, unplaced, n)
    return psol

def rand_gradient_projection(problem, n = 10):
    '''
    Call the gradient projection heuristic n times with n random
    initial solutions.

    Should provide a better solution than the regular gradient_projection.
    Note: Testing of this function is limited since it takes a long time to
    execute.

    :param problem: the problem from which to generate the partial solution
    :type problem: Problem
    :param n: the number of initial solutions to try
    :type n: int
    
    Examples:
    
    >>> from pyQAP import *
    >>> vars.verbose = False
    >>> prob = problems.QAPLIB('nug12')
    >>> rand_gradient_projection(prob,5)
    11 Improvements. Overall 15.507714
    11 Improvements. Overall 26.188623
    11 Improvements. Overall 26.534933
    11 Improvements. Overall 54.816300
    11 Improvements. Overall 18.372427
    Partial QAP solution: [ 9  1  7  4 11  6  8  3  5 10  0  2] with 12 placed.
     Placed:12
     Unplaced:[0 0 0 0 0 0 0 0 0 0 0 0]
    >>> prob = problems.QAPLIB('had16')
    >>> rand_gradient_projection(prob,5)
    11 Improvements. Overall 18.387210
    11 Improvements. Overall 24.930627
    11 Improvements. Overall 17.901750
    11 Improvements. Overall 22.392300
    11 Improvements. Overall 18.745858
    Partial QAP solution: [15  1  9  2  3  5  4 11  0 12 10 14  8 13  6  7] with 16 placed.
     Placed:16
     Unplaced:[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
    '''
    best_psol = None
    best_cost = float('inf')
    for i in xrange(n):
        new_perm = hh.rand_perm(problem.n)
        new_psol = gradient_projection(problem, initial = new_perm)
        new_cost = new_psol.cost(problem)
        if new_cost < best_cost:
            best_psol = new_psol
            best_cost = new_cost
    return best_psol

def simulated_annealing(problem, mult = 2, steps =1000,\
                        initial = None):
    '''
    Return a heuristic psol produced by simulated annealing for a given number
    of iterations.

    t_0 is the initial temperature and the alpha value controls
    the exponential decay of the temperature. Initial is the optional argument
    for a starting solution.

    :param problem: the problem from which to generate the partial solution
    :type problem: Problem
    :param mult: the constant to multiply by n to get the number of steps
    :type mult: int
    
    Examples:
    
    >>> from pyQAP import *
    >>> prob = problems.QAPLIB('nug12')
    >>> simulated_annealing(prob)
    Cost: 670.000000
    Partial QAP solution: [11  3  2  4  8  0  9  1  6  7 10  5] with 12 placed.
     Placed:12
     Unplaced:[0 0 0 0 0 0 0 0 0 0 0 0]
    >>> prob = problems.QAPLIB('sko100f')
    >>> simulated_annealing(prob)
    Cost: 151506.000000
    Partial QAP solution: [42  1 78 81 88 60 15  2 56 79 62 51 66  5  4 28 68 37 84 96 38 92 59  7 86
     70 65 73 11 99 21 94 48  8 35  6 67 25 12 71 72  0 63 17 83 75 90 55 95 22
     16 87 36 23 45 98 80 76 40 53 54 69 34 24 74 64 50 82 29 57 85 97 10 14 13
     49 46 39 41 27  9 58 18 52 19  3 47 33 26 32 20 91 77 61 93 31 43 44 89 30] with 100 placed.
     Placed:100
     Unplaced:[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
     0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
     0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
    >>> prob = problems.QAPLIB('had16')
    >>> simulated_annealing(prob)
    Cost: 3738.000000
    Partial QAP solution: [ 2 12  4  1 11 13 10  7  9  5  6 14  3  8 15  0] with 16 placed.
     Placed:16
     Unplaced:[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
    '''
    #determine parameters
    n = problem.n
    t = n*problem.f.mean()*problem.d.mean()
    if t < 1:
        t = 1000
    alpha = exp(-log(t)/steps)
    #generate initial psol if nessecary
    if initial == None:
        psol = classes.PSolution()
        psol.n = n
        psol.sol = np.random.permutation(n)
        psol.unplaced = np.zeros((n), dtype=np.int)
        psol.placed = n
    else:
        psol = initial
    #main loop
    sol = psol.sol
    frozen = False
    iterations = n*mult
    while (not frozen):
        changes = 0
        for i in xrange(iterations):
            #compute delta
            i1 = np.random.randint(n)
            i2 = np.random.randint(n)
            delta = cost.get_delta(i1, i2, sol, problem.f, problem.d, problem.n)
            #see if change made
            if delta < 0.0:
                psol.sol[i1], psol.sol[i2] = psol.sol[i2], psol.sol[i1]
                changes += 1
            else:
                #if p(delta) < rand() accept and swap
                if np.random.random() < exp(-delta/t):
                    #swap
                    psol.sol[i1], psol.sol[i2] = psol.sol[i2], psol.sol[i1]
                    changes += 1
        #print "Temp: %f Changes: %d"%(t, changes)
        frozen = not changes
        t = t*alpha
        new_cost = psol.cost(problem)
        if t < 1e-10:
            break
    if vars.verbose:
        print "Cost: %f"%(new_cost)
    return psol

def parallel_annealing(problem, n = None, threads = vars.threads,  mult = 2,  steps = 1000):
    '''
    Performs simulated annealing in parallel n times.

    See simulated_annealing for more info. Returns all of the
    generated partial solutions.

    :param problem: the problem from which to generate the partial solution
    :type problem: Problem
    :param n: the number of independent annealings
    :type n: int
    :param threads: the number of threads to use
    :type threads: int
    :param mult: the constant to multiply by n to get the number of steps
    :type mult: int
    :param steps: the approximate number of steps
    :param steps: int
    
    Examples:
    
    >>> from pyQAP import *
    >>> vars.verbose = False
    >>> prob = problems.QAPLIB('nug12')
    >>> parallel_annealing(prob,10)
    Best Cost: ...
    Partial QAP solution: [...] with 12 placed.
     Placed:12
     Unplaced:[0 0 0 0 0 0 0 0 0 0 0 0]
    >>> prob = problems.QAPLIB('sko100f')
    >>> parallel_annealing(prob,10)
    Best Cost: ...
    Partial QAP solution: [...] with 100 placed.
     Placed:100
     Unplaced:[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
     0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
     0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
    >>> prob = problems.QAPLIB('had16')
    >>> parallel_annealing(prob,10)
    Best Cost: ...
    Partial QAP solution: [...] with 16 placed.
     Placed:16
     Unplaced:[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
    '''
    if n == None:
        n = problem.n*10
    psols = hh.parallel_annealing_all(n, threads, problem, mult,  steps)
    get_cost = lambda x: x.cost(problem)
    best = min(psols, key = get_cost)
    if vars.verbose:
        print "Best Cost: %f"%(best.cost(problem))
    return best
